package cn.czk.demo;

import cn.czk.demo.utils.MyDocumentSplitter;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.zhipu.ZhipuAiEmbeddingModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;

import java.net.URISyntaxException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.List;

public class MeituanRagLoader {

    public static void main(String[] args) throws URISyntaxException {
        //加载文件并封装成 Document
        Path documentPath = Paths.get(MeituanRagLoader.class.getClassLoader().getResource("meituanrag.txt").toURI());
        TextDocumentParser documentParser = new TextDocumentParser();
        Document document = FileSystemDocumentLoader.loadDocument(documentPath, documentParser);

        //切分文件
        MyDocumentSplitter splitter = new MyDocumentSplitter();
        List<TextSegment> segments = splitter.split(document);

        //文本向量化
        ZhipuAiEmbeddingModel model = ZhipuAiEmbeddingModel.builder()
                .apiKey("53e452feb1164ec8a6c9cd9112a37d60.xTOwz9t4NIZuYTYb")
                .logRequests(true)
                .logResponses(true)
                .maxRetries(1)
                .callTimeout(Duration.ofSeconds(60))
                .connectTimeout(Duration.ofSeconds(60))
                .writeTimeout(Duration.ofSeconds(60))
                .readTimeout(Duration.ofSeconds(60))
                .build();

        List<Embedding> embeddings = model.embedAll(segments).content();

        //向量数据库保存向量化结果
        EmbeddingStore<TextSegment> embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(512)
                .indexName("meituan-rag")
                .build();

        embeddingStore.addAll(embeddings, segments);

        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore) //向量存储模型
                .embeddingModel(model) //向量模型
                .maxResults(5) // 最相似的5个结果
                .minScore(0.8) // 只找相似度在0.8以上的内容
                .build();

        String question = "在线⽀付取消订单后钱怎么返还？"; //⽤户的问题
        Query query = new Query(question);
        List<Content> contentList = contentRetriever.retrieve(query);
        for (Content content : contentList) {
            System.out.println(content);
        }

        //检索相关信息

        //构建 Prompt 提示词

    }


}
